package com.mti.digital.security.code.generation;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.AutoGenerator;
import com.baomidou.mybatisplus.generator.InjectionConfig;
import com.baomidou.mybatisplus.generator.config.*;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.config.rules.DateType;
import com.baomidou.mybatisplus.generator.config.rules.NamingStrategy;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;
import lombok.Builder;
import lombok.Data;

import java.io.File;
import java.nio.file.Paths;
import java.util.*;

/**
 * @ClassName:
 * @description: 代码生成工具
 * @author: Allen
 * @create: 2020/8/31 16:52
 **/
@Data
@Builder
public class CodeGeneratorTool {

    public static void main(String[] args) {
        CodeGeneratorTool.builder()
                .driverName("com.mysql.cj.jdbc.Driver")
                .url("jdbc:mysql://127.0.0.1:3306/master?serverTimezone=UTC&useSSL=true")
                .username("sonar").password("123456")
                .outputDir("D://tmp")//输出目录,可选，默认本项目录s
                .projectName("digital-security-parent")//项目目录,可选，默认不包含项目目录
                .tableInclude(new String[]{"t_order"})
                .packageName("com.mti.digital.security.dispose")//包名
                .tablePrefix("auth,t_,r_,d_,t_pf_,t_gx_")//前缀
                .build().autoGenerator();
    }


    private String outputDir;

    /***
     * 数据库连接
     */
    private String url;

    /***
     * 驱动名称
     */
    private String driverName;

    /***
     * 用户名
     */
    private String username;

    /***
     * 密码
     */
    private String password;

    /***
     * 模式
     */
    private String schemaName;

    /***
     * 作者
     */
    private String author;

    /***
     * 项目名称
     */
    private String projectName;

    /***
     * 包名称
     */
    private String packageName;

    /***
     * 表前缀
     */
    private String tablePrefix;

    /***
     * 自定义待生成表
     */
    private String[] tableInclude;

    private static final String MAIN_PATH = "/src/main/java";
    private static final String MAIN_RESOURCES_PATH = "/src/main/resources";
    private static final String TEST_PATH = "/src/test/java";
    private static final String TEST_RESOURCES_PATH = "/src/test/resources";



    public void autoGenerator() {
        // 代码生成器
        AutoGenerator mpg = new AutoGenerator();

        // 全局配置
        GlobalConfig globalConfig = new GlobalConfig();

        if (StringUtils.isNotBlank(outputDir)) {
            globalConfig.setOutputDir(outputDir + MAIN_PATH);
            globalConfig.setOutputDir(outputDir + File.separator + Optional.ofNullable(projectName).orElse("") + MAIN_PATH);
        } else {
            String projectPath = System.getProperty("user.dir");
            globalConfig.setOutputDir(projectPath + MAIN_PATH);
        }

        if (author != null && author.length() > 0) {
            globalConfig.setAuthor(author);
        } else {
            globalConfig.setAuthor("System");
        }
        globalConfig.setDateType(DateType.ONLY_DATE);
        globalConfig.setOpen(false);
        globalConfig.setFileOverride(false); //重新生成时文件是否覆盖
        globalConfig.setDateType(DateType.TIME_PACK);
        globalConfig.setSwagger2(true); //实体属性 Swagger2 注解
        mpg.setGlobalConfig(globalConfig);

        // 数据源配置
        DataSourceConfig dataSourceConfig = new DataSourceConfig();
        dataSourceConfig.setUrl(url);
        if (schemaName == null || schemaName.length() == 0) {
            dataSourceConfig.setSchemaName("public");
        } else {
            dataSourceConfig.setSchemaName(schemaName);
        }
        dataSourceConfig.setDriverName(driverName);
        dataSourceConfig.setUsername(username);
        dataSourceConfig.setPassword(password);
        mpg.setDataSource(dataSourceConfig);

        // 包配置
        PackageConfig packageConfig = new PackageConfig();
        packageConfig.setParent("com.baomidou.ant");
        if (packageName != null && packageName.length() > 0) {
            packageConfig.setParent(packageName);
        }
        packageConfig.setEntity("model.entity");
        mpg.setPackageInfo(packageConfig);

        // 自定义属性注入
        InjectionConfig injectionConfig = new InjectionConfig() {
            @Override
            public void initMap() {
                Map<String, Object> map = new HashMap<>();
                map.put("driverName", driverName);
                map.put("url", url);
                map.put("username", username);
                map.put("password", password);
                map.put("packageName", packageName);
                map.put("projectName", projectName);
                this.setMap(map);
            }
        };

        // 模板输出
        myTemplateOut(injectionConfig);
        // 配置自定义属性注入
        mpg.setCfg(injectionConfig);

        // 配置模板
        TemplateConfig templateConfig = new TemplateConfig();
        mpg.setTemplate(templateConfig);

        // 策略配置
        StrategyConfig strategy = new StrategyConfig();
        strategy.setNaming(NamingStrategy.underline_to_camel);
        strategy.setColumnNaming(NamingStrategy.underline_to_camel);
        strategy.setEntityLombokModel(true);
        strategy.setRestControllerStyle(true);

        strategy.setControllerMappingHyphenStyle(true);
        strategy.setTablePrefix(packageConfig.getModuleName() + "_");
        if (tablePrefix == null || tablePrefix.length() == 0) {
            strategy.setTablePrefix("s_", "t_", "bo_", "r_", "d_");
        } else {
            strategy.setTablePrefix(tablePrefix.split(","));
        }
        if (tableInclude != null && tableInclude.length > 0) {
            strategy.setInclude(tableInclude);
        }

        mpg.setStrategy(strategy);
        // 设置模板引擎
        mpg.setTemplateEngine(new FreemarkerTemplateEngine());
        mpg.execute();
    }

    private void myTemplateOut(InjectionConfig cfg) {
        // 自定义输出配置
        List<FileOutConfig> focList = new ArrayList<>();
        // 模板路径
        String templatesPath = this.getClass().getClassLoader().getResource("templates").getPath();

        List<File> files = new LinkedList<>();
        traverseFolder2(new File(templatesPath), files);

        if (files.isEmpty()) {
            return;
        }

        for (File file : files) {

            String fileRname = file.getName();
            StringBuilder outFileName = new StringBuilder(outputDir + "/" + projectName);
            String path = file.getAbsolutePath();
            if (Paths.get(path).toUri().toString().contains("main/")) {
                outFileName.append((path.contains(".java") ? MAIN_PATH : MAIN_RESOURCES_PATH) + "/");
                fileRname = "/main/" + file.getName();
            } else if (Paths.get(path).toUri().toString().contains("test/")) {
                outFileName.append((path.contains(".java") ? TEST_PATH : TEST_RESOURCES_PATH) + "/");
                fileRname = "/test/" + file.getName();
            } else {
                outFileName.append("/");
            }

            String templatePath = "/templates/" + fileRname;
            focList.add(new FileOutConfig(templatePath) {
                @Override
                public String outputFile(TableInfo tableInfo) {
                    return outFileName.toString() + file.getName().replace("${pn}", packageName.replace(".", "/")).replace("${entity}", tableInfo.getEntityName()).replace("_", "/").replace(".ftl", "");

                }
            });

        }

        cfg.setFileOutConfigList(focList);
    }

    /***
     * 查询文件
     * @param path
     */
    public static List<File> traverseFolder2(File path, List<File> filesList) {
        if (!path.exists()) {
            return new ArrayList<>();
        }
        File[] files = path.listFiles();
        if (null == files || files.length == 0) {
            return new ArrayList<>();
        }
        for (File file2 : files) {
            if (file2.isDirectory()) {
                traverseFolder2(file2, filesList);
            } else {
                filesList.add(file2);
            }
        }
        return filesList;
    }
}
